from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import Crypto.Hash.TupleHash128
# CryptSM4是SM4密码算法的一种实现，SM4_ENCRYPT和SM4_DECRYPT分别表示SM4加密和SM4解密的方法。

# num_set = [chr(i) for i in range(48, 58)]
# char_set = [chr(i) for i in range(97, 123)]
# total_set = num_set + char_set
#这里用于生成英文和数字字符，合成为所有字符的set
crypt_sm4 =CryptSM4()

#获取key,key会进行一次hash
def get_key(key):
    hd= Crypto.Hash.TupleHash128.new(digest_bytes=16)
    key=hd.update(key.encode("UTF-8")).hexdigest()
    half=len(key) // 2
    # print("key:",key)
    return key[:half]


#获取/生成随机iv,cbc模式下需要初始向量IV,如果不想使用iv也可以使用ebc模式，但是相对不安全
#这里决定用hash将key分为32位，前16位用于key，后16位用于iv
def get_iv(iv):
    hd= Crypto.Hash.TupleHash128.new(digest_bytes=16)
    iv=hd.update(iv.encode("UTF-8")).hexdigest()
    half=len(iv) // 2
    # print("iv:",iv)
    return iv[half:]

#编写SM4_cbc模式的加密函数SM4_cbc_enc
def SM4_cbc_enc(plaintext,iv,key):
    crypt_sm4.set_key(key,SM4_ENCRYPT)
    ciphertext_bytes=crypt_sm4.crypt_cbc(iv,plaintext)
    return ciphertext_bytes

#编写SM4_cbc模式的解密函数SM4_cbc_dec
def SM4_cbc_dec(ciphertext,iv,key):
    crypt_sm4.set_key(key,SM4_DECRYPT)
    plaintext_bytes=crypt_sm4.crypt_cbc(iv,ciphertext)
    return plaintext_bytes

#输入文件转化为bytes
def exp_file(filename):
    with open(filename,"rb") as f:
        plaintext_bytes= f.read()
    return plaintext_bytes




#key可以自定义
# key="1234123412341234"
# key=get_key(key,16)
# key_bytes=bytes(key,encoding="utf8")
# with open('D:/key_bytes.txt','wb') as f:
#     f.truncate()
#     f.write(key_bytes)
# f.close
# print("key:",key_bytes)

# iv=gen_iv(16)
# iv_bytes=bytes(iv,encoding="utf8")
# with open('D:/iv_bytes.txt','wb')as f:
#     f.truncate
#     f.write(iv_bytes)
# f.close
# print("iv:",iv_bytes)

# plaintext_bytes=exp_file(input())
# ciphertext_bytes=SM4_cbc_enc(plaintext_bytes,iv_bytes,key_bytes)
# print("密文:",ciphertext_bytes.hex())
#写入文件
# plaintext=SM4_cbc_dec(ciphertext_bytes,iv_bytes,key_bytes)
# print("明文:",plaintext)






